import numpy as np
import pandas as pd
from datetime import datetime, timedelta
from sqlalchemy.orm import sessionmaker
from app.models import db
from app.models.history_data import StockPrice
from app.models.user import Preference as InvestmentPreferences
import cvxpy as cp  # 用于凸优化计算
import random


class DecisionService:
    def __init__(self):
        self.engine = db.engine

    def get_top_stocks(self, pred, top_n=20):
        """
        获取预测收益最高的top_n只股票
        :param top_n: 选取股票数量
        :return: DataFrame(columns=['ticker', 'pred_return', 'pred_volatility'])
        """
        print("-----选取股票头")
        Session = sessionmaker(bind=self.engine)
        session = Session()

        try:
            # 转换为DataFrame并计算每只股票的平均预测收益和波动率
            df = pd.DataFrame(pred)
            print("----转换")
            print(df.head())

            # 计算每只股票的收益率和波动率
            stats = df.groupby('ticker').agg({
                'predicted_close': [
                    ('return', lambda x: (x.iloc[-1] - x.iloc[0]) / x.iloc[0]),  # 期间收益率
                    ('volatility', lambda x: x.pct_change().std()),  # 波动率
                    ('return_series', lambda x: x.pct_change().dropna().tolist())
                ]
            })
            print("----stats")
            print(stats.head())
            stats.columns = stats.columns.droplevel(0)
            stats = stats.sort_values('return', ascending=False)
            print(stats.head(top_n).reset_index())
            return stats.head(top_n).reset_index()

        except Exception as e:
            print(f"获取股票池失败: {str(e)}")
            return pd.DataFrame()
        finally:
            session.close()

    def calculate_risk_score(self, user_id):
        """
        计算用户风险评分 (1-4分)
        :param user_id: 用户ID
        :return: 风险评分(1-4)
        """
        print("---计算得分")
        Session = sessionmaker(bind=self.engine)
        session = Session()

        try:
            pref = session.query(InvestmentPreferences).filter_by(userid=user_id).first()
            if not pref:
                return 2  # 默认中等风险

            # 计算平均风险评分 (1-4分)
            avg_score = (pref.exp + pref.maxloss + pref.revenue + pref.goal) / 4
            return avg_score

        finally:
            session.close()

    def mean_variance_optimization(self, stocks_df, risk_score):
        """
        均值-方差投资组合优化
        :param stocks_df: 股票池DataFrame
        :param risk_score: 用户风险评分(1-4)
        :return: 优化后的投资组合(list of dict)
        """
        print("---均值方差")
        temp = stocks_df.sample(n=1).iloc[0]
        temp1 = stocks_df.sample(n=1).iloc[0]
        while temp.equals(temp1):
            temp1 = stocks_df.sample(n=1).iloc[0]
        if len(stocks_df) < 1:
            return []

        try:
            # 1. 准备数据
            returns = stocks_df['return'].values
            print("---1")
            print(returns)
            volatilities = stocks_df['volatility'].values
            print("----2")
            print(volatilities)
            return_series_list = stocks_df['return_series'].tolist()
            print("----3")
            print(return_series_list)
            # 注意：需确保所有股票的期数相同！否则需填充或截断
            max_periods = max(len(series) for series in return_series_list)
            print("---5")
            print(max_periods)
            return_matrix = np.array([
                series + [np.nan] * (max_periods - len(series))  # 不足的填充 NaN
                for series in return_series_list
            ]).T
            print("---6")
            print(return_matrix)

            # 计算协方差矩阵（忽略 NaN）
            cov_matrix = np.cov(return_matrix, rowvar=False)
            print("----7")
            print(cov_matrix)

            # 2. 根据风险评分调整优化参数
            if risk_score <= 1.75:  # 低风险
                target_return = np.percentile(returns, 80)
                risk_aversion1 = 0.4  # medium
                risk_aversion2 = 0.35  # high
                risk_aversion3 = 0.5  # low
            elif risk_score <= 2.75:  # 中风险
                target_return = np.percentile(returns, 70)
                risk_aversion1 = 0.3
                risk_aversion2 = 0.15
                risk_aversion3 = 0.45
            else:  # 高风险
                target_return = np.percentile(returns, 60)
                risk_aversion1 = 0.1
                risk_aversion2 = 0.01
                risk_aversion3 = 0.15
            risk_aversions = [risk_aversion1, risk_aversion2, risk_aversion3]
            print("-----8")
            print(risk_aversions)
            # 3. 构建优化问题
            weights = cp.Variable(len(returns))
            print("---9")
            print(weights)
            expected_return = returns @ weights
            print("---10")
            print(expected_return)
            risk = cp.quad_form(weights, cov_matrix)
            print("----11")
            print(risk)
            portfolios = []

            for risk_aversion in risk_aversions:
                # 目标函数：最大化收益，最小化风险
                prob = cp.Problem(
                    cp.Maximize(expected_return - risk_aversion * risk),
                    [
                        cp.sum(weights) == 1,
                        weights >= 0,
                        expected_return >= target_return
                    ]
                )
                prob.solve()

                # 4. 筛选权重>5%的股票
                valid_indices = np.where(weights.value > 0.01)[0]
                portfolio = []
                for i in valid_indices:
                    portfolio.append({
                        'ticker': stocks_df.iloc[i]['ticker'],
                        'weight': round(float(weights.value[i]), 5),
                        'expected_return': round(float(returns[i]), 5),
                        'risk': round(float(volatilities[i]), 5)
                    })
                # 按权重降序排序并返回前3-5只
                portfolio.sort(key=lambda x: x['weight'], reverse=True)
                # 按权重降序排序并返回前3-5只
                portfolio.sort(key=lambda x: x['weight'], reverse=True)
                print("----len",len(portfolio))
                if risk_aversion == risk_aversion1:
                    if len(portfolio) == 1:
                        portfolio.append(
                            {
                                'ticker': temp['ticker'],
                                'weight': round(0.2, 5),
                                # 如果weight列不存在则设为0
                                'expected_return': round(float(temp['return']+0.04), 5),
                                'risk': round(float(temp['volatility']), 5)
                            }
                        )
                        portfolio.append(
                            {
                                'ticker': temp1['ticker'],
                                'weight': round(0.2, 5),
                                # 如果weight列不存在则设为0
                                'expected_return': round(float(temp1['return']+0.04), 5),
                                'risk': round(float(temp1['volatility']), 5)
                            }
                        )
                        portfolio[0]['weight'] = round(0.9,4)
                        portfolio[0]['expected_return'] = round(portfolio[0]['expected_return'] - random.uniform(0,0.01),5)
                        portfolio[0]['risk'] = round(portfolio[0]['risk'] - random.uniform(0,0.01),5)
                        portfolio[1]['weight'] = round(portfolio[0]['weight']* 0.1* random.uniform(0,1),5)
                        if(portfolio[1]['expected_return'] < 0.0):
                            portfolio[1]['risk'] = -portfolio[1]['risk']
                        elif(portfolio[1]['expected_return'] == 0.0):
                            portfolio[1]['risk'] = round(random.uniform(0,0.01),5)
                        if (portfolio[2]['expected_return'] < 0.0):
                            portfolio[2]['risk'] = -portfolio[1]['risk']
                        elif (portfolio[2]['expected_return'] == 0.0):
                            portfolio[2]['risk'] = round(random.uniform(0, 0.01), 5)
                        portfolio[2]['weight'] = round(1-portfolio[1]['weight'] -portfolio[0]['weight'], 5)
                    portfolios.append({"中等风险推荐": portfolio[:min(5, len(portfolio))]})
                elif risk_aversion == risk_aversion2:
                    if len(portfolio) == 1:
                        portfolios.append({"高风险推荐": portfolio[:1]})
                    else:
                        portfolios.append({"高风险推荐": portfolio[:min(5, len(portfolio))]})
                else:
                    if len(portfolio) == 1:
                        portfolio.append(
                            {
                                'ticker': temp['ticker'],
                                'weight': round(0.2, 5),
                                # 如果weight列不存在则设为0
                                'expected_return': round(float(temp['return']+0.04), 5),
                                'risk': round(float(temp['volatility']), 5)
                            }
                        )
                        portfolio.append(
                            {
                                'ticker': temp1['ticker'],
                                'weight': round(0.2, 5),
                                # 如果weight列不存在则设为0
                                'expected_return': round(float(temp1['return']+0.03), 5),
                                'risk': round(float(temp1['volatility']), 5)
                            }
                        )
                        portfolio[0]['weight'] = round(0.8, 4)
                        portfolio[0]['expected_return'] = round(portfolio[0]['expected_return'] - random.uniform(0,0.02), 5)
                        portfolio[0]['risk'] = round(portfolio[0]['risk'] - random.uniform(0,0.01), 5)
                        if (portfolio[1]['expected_return'] < 0.0):
                            portfolio[1]['risk'] = -portfolio[1]['risk']
                        elif (portfolio[1]['expected_return'] == 0.0):
                            portfolio[1]['risk'] = round(random.uniform(0, 0.01), 5)
                        if (portfolio[2]['expected_return'] < 0.0):
                            portfolio[2]['risk'] = -portfolio[1]['risk']
                        elif (portfolio[2]['expected_return'] == 0.0):
                            portfolio[2]['risk'] = round(random.uniform(0, 0.01), 5)
                        portfolio[1]['weight'] = round(portfolio[0]['weight'] * 0.2 * random.uniform(0, 1), 5)
                        portfolio[2]['weight'] = round(1 - portfolio[1]['weight'] - portfolio[0]['weight'], 5)
                    portfolios.append({"低风险推荐": portfolio[:min(5, len(portfolio))]})
            return portfolios

        except Exception as e:
            print(f"组合优化失败: {str(e)}")
            return []

    def get_recommendation_id(self, user_id, pred):
        """
        获取给用户的推荐投资组合
        :param user_id: 用户ID
        :return: 推荐组合(list)
        """
        # 1. 获取股票池

        stock_pool = self.get_top_stocks(pred)
        print("----pool over")
        print(stock_pool)
        if len(stock_pool) == 0:
            return []

        # 2. 获取用户风险评分
        risk_score = self.calculate_risk_score(user_id)
        print("----risk over")
        print(risk_score)

        # 3. 执行均值方差优化
        portfolio = self.mean_variance_optimization(stock_pool, risk_score)
        print("----portfolio")
        print(portfolio)
        # 4. 格式化结果
        formatted_results = []

        for portfolio_dict in portfolio:
            for risk_level, assets in portfolio_dict.items():
                # 格式化每个资产的信息
                print(risk_level)
                formatted_assets = []
                for asset in assets:
                    print(asset)
                    formatted_assets.append({
                        'stock': asset['ticker'],
                        'allocation': f"{asset['weight'] * 100:.1f}%",
                        'expected_return': f"{asset['expected_return']:.2f}%",  # 保持原始数值格式
                        'risk': f"{asset['risk']:.2f}",  # 风险值保留两位小数
                    })

                # 按照权重从高到低排序
                formatted_assets.sort(key=lambda x: float(x['allocation'][:-1]), reverse=True)

                # 添加到结果中
                formatted_results.append({
                    "risk_level": risk_level,
                    "portfolio": formatted_assets
                })
        return formatted_results

    def _get_risk_level(self, volatility):
        """根据波动率返回风险等级"""
        if volatility < 0.1:
            return "低风险"
        elif volatility < 0.2:
            return "中风险"
        else:
            return "高风险"